import tensorflow.compat.v1 as tf
import tensorflow as tsf
import numpy as np
import os

# 使用tensorflow框架，利用循环神经网络LSTM算法进行短句训练“ how old are you”。
x_sentence = "how old are you"

# (一)	建立字典（8分）
letters_set = set(list(x_sentence))
dict_len = len(letters_set)
idx2char = list(letters_set)
char2idx = {}
for i, ch in enumerate(idx2char):
    char2idx[ch] = i

# (二)	构造数据集x_data,y_data（8分）
ori_sentence = x_sentence
x_sentence = ori_sentence[:-1]
x_idx = [char2idx[ch] for ch in x_sentence]
x_data = np.eye(dict_len)[x_idx]
y_sentence= ori_sentence[1:]
y_data = [char2idx[ch] for ch in y_sentence]
x_data = np.array([x_data])
y_data = np.array([y_data])
print(x_sentence)
print(y_sentence)
print(np.shape(x_data))
print(np.shape(y_data))

# (三)	设置参数（8分）
ver = 'v1.3'
learning_rate = 0.01
n_iters = 200
n_hidden = 20
n_layers = 4
n_steps = len(x_sentence)

# (四)	定义占位符（8分）
with tf.variable_scope('Input'):
    ph_x = tf.placeholder(tf.float32, [None, n_steps, dict_len], 'ph_x')
    ph_y = tf.placeholder(tf.int32, [None, n_steps], 'ph_y')
    n_samples = tf.shape(ph_x)[0]

# (五)	定义LSTM单元（4分）、堆叠多层RNN单元（4分）、调用动态RNN函数（4分）
with tf.variable_scope('RNN'):
    cell_arr = [tf.nn.rnn_cell.LSTMCell(n_hidden) for i in range(n_layers)]
    cell = tf.nn.rnn_cell.MultiRNNCell(cell_arr)
    outputs, states = tf.nn.dynamic_rnn(cell, ph_x, dtype=tf.float32)

# (六)	定义全连接层（8分）
with tf.variable_scope('FC'):
    outputs = tf.reshape(outputs, [-1, n_hidden])
    fc = tsf.contrib.layers.fully_connected(outputs, dict_len, activation_fn=None)
    logits = tf.reshape(fc, [-1, n_steps, dict_len])

# (七)	计算序列损失（8分）
with tf.variable_scope('Cost'):
    weights = tf.ones([n_samples, n_steps], dtype=tf.float32, name='weights')
    cost = tf.reduce_mean(tsf.contrib.seq2seq.sequence_loss(logits=logits, targets=ph_y, weights=weights))
    tf.summary.scalar('cost', cost)

# train
with tf.variable_scope('Train'):
    train = tf.train.AdamOptimizer(learning_rate=learning_rate).minimize(cost)

# (八)	定义准确率计算模型（8分）
with tf.variable_scope('Metrics'):
    predict = tf.argmax(logits, axis=2)
    acc = tf.reduce_mean(
        tf.cast(
            tf.equal(
                tf.cast(predict, dtype=tf.int32),
                tf.cast(ph_y, dtype=tf.int32)
            ),
            tf.float32
        )
    )
    tf.summary.scalar('acc', acc)

# summary
with tf.variable_scope('Summary'):
    summary = tf.summary.merge_all()

with tf.Session() as sess:
    saver = tf.train.Saver()
    save_path = './_save/' + os.path.basename(__file__) + '_' + ver
    if os.path.exists(save_path + '.meta'):
        saver.restore(sess, save_path)
        print('LOADED SAVED SESSION.')
    else:
        print('Start to train ...')
        with tf.summary.FileWriter('./_log/' + os.path.basename(__file__), sess.graph) as fw:
            sess.run(tf.global_variables_initializer())

            # (九)	训练迭代100次（8分）
            for i in range(n_iters):
                _, costv, accv, sv = sess.run([train, cost, acc, summary], feed_dict={ph_x: x_data, ph_y: y_data})
                fw.add_summary(sv, i)

                # (十)	输出损失值、准确率（8分）
                print(f'#{i + 1}: cost = {costv}, acc = {accv}')
                fw.flush()
                if np.isclose(1.0, accv):
                    break

        saver.save(sess, save_path)
        print('Training over. Saved session.')

    # (十一)	预测结果查字典后输出字符串（8分）
    h_idx = sess.run(predict, feed_dict={ph_x: x_data})
    h_sentence = [''.join([idx2char[i] for i in row]) for row in h_idx]
    print('预测结果查字典后输出字符串:')
    print(h_sentence)

    # (十二)	用一个新的数据“how old you are”进行测试，以字符串格式输出预测的结果（8分）
    # new_string_arr = ["how old you are"]
    new_string_arr = [
        "how old are you",
        "how old you are",
        "you are how old",
        "how you old are"
    ]
    print(f'用一个新的数据“{new_string_arr}”进行测试，以字符串格式输出预测的结果:')
    x_idx = np.array([[char2idx[ch] for ch in new_string[:-1]] for new_string in new_string_arr], dtype=np.int32)
    x_data = np.eye(dict_len)[x_idx]
    h_idx = sess.run(predict, feed_dict={ph_x: x_data})
    h_sentence = [''.join([idx2char[i] for i in row]) for row in h_idx]
    print(h_sentence)
